from typing import Dict
import pandas as pd
import numpy as np
from nicegui import ui
from functools import cache

# region df 辅助方法


def continue_code(x_series: pd.Series):
    """
    为每个连续区域打上唯一编号
    比如: ['a','a','b','a','c','c']
    结果：[1,1,2,3,4,4]
    """
    return (x_series.shift(1) != x_series).cumsum()


def calc_span(df: pd.DataFrame, target_col: str):
    continue_key = continue_code(df[target_col])
    is_first_cell = df[target_col].shift() != df[target_col]
    gp_len = df.groupby(continue_key)[target_col].transform(len)
    return np.where(is_first_cell, gp_len, 1)


# endregion


# region aggrid 辅助方法
class AggridOptionsUtils:
    def __init__(self, options: Dict) -> None:
        self._options = options

    @cache
    def __get_column_data_map(self):
        if "columnDefs" not in self._options:
            return {}

        return {col["field"]: col for col in self._options["columnDefs"]}

    def __get_column_data(self, column: str) -> Dict:
        return self.__get_column_data_map()[column]

    @staticmethod
    def from_dataframe(df: pd.DataFrame):
        return AggridOptionsUtils(
            {
                "columnDefs": [{"field": str(col)} for col in df.columns],
                "rowData": df.to_dict("records"),
            }
        )

    def set_column_value(self, column: str, name: str, value):
        data = self.__get_column_data(column)
        if isinstance(value, dict):
            if name not in data:
                data[name] = {}
            data[name].update(value)
        else:
            data[name] = value

        return self

    def hide_column(self, column: str):
        return self.set_column_value(column, "hide", True)

    def row_span(self, column: str, span_column: str):
        self._options["suppressRowTransform"] = True
        return self.set_column_value(
            column,
            ":rowSpan",
            rf"""params => {{
              return params.data['{span_column}']
            }}""",
        ).set_column_value(
            column,
            "cellClassRules",
            {
                ":cell-span": rf"""params => {{
              return params.data['{span_column}']!==1
            }}"""
            },
        )


# endregion


ui.add_head_html(
    """
<style>
.cell-span.ag-cell-value{
    background-color: white;
    border-left:1px solid #d8dbde !important;
    border-bottom:1px solid #d8dbde !important;
    border-right:1px solid #d8dbde !important;
    display:flex !important;
    place-items:center;
}

/*选中合并单元格时的样式*/
.cell-span.ag-cell-focus{
    background-color: red;
}

/*鼠标在所在行时的合并单元格样式*/
.ag-row-hover > .cell-span{
   /* background-color: green;*/
}

</style>              
"""
)

df = pd.DataFrame(
    {
        "name": list("aaabbcaa"),
        "other": list("xyzxyzzz"),
        "value1": range(8),
        "value2": range(8),
    }
)

df["_span"] = calc_span(df, "name")


grid_utils = (
    AggridOptionsUtils.from_dataframe(df)
    .row_span("name", "_span")
    .row_span("value1", "_span")
    .hide_column("_span")
)


ui.aggrid(grid_utils._options)


ui.run()
